import numpy as np

import matplotlib.pyplot as plt

from util.util_ris_pattern import phaseBit_2_deg, phase_2_pattern_xyz_fft, point_2_phase, phase_2_bit

from multi_beam_trace.util_analysis import compare_arrays

from multi_beam.multi_beam_PS_complex import psm_complex_beam_n


# ================================================== 生成轨迹 =================================================
def create_scan_traces_theta_phi_beam_2(theta1_start, theta1_step, theta1_end, phi1_start, phi1_step, phi1_end,
                                        theta2_start, theta2_step, theta2_end, phi2_start, phi2_step, phi2_end):
    traces_list = []
    if theta1_step > 0:
        length = (theta1_end - theta1_start) // theta1_step
    elif phi1_step > 0:
        length = (phi1_end - phi1_start) // phi1_step
    elif theta2_step > 0:
        length = (theta2_end - theta2_start) // theta2_step
    elif phi2_step > 0:
        length = (phi2_end - phi2_start) // phi2_step
    else:
        length = 0
    for i in range(0, length, 1):
        theta1 = theta1_start + i * theta1_step
        phi1 = phi1_start + i * phi1_step
        theta2 = theta2_start + i * theta2_step
        phi2 = phi2_start + i * phi2_step
        traces_list.append([[theta1, phi1], [theta2, phi2]])
    return traces_list

# ================================================== 方向图相关 =========================================================
def trace_2_phaseBit_PS_complex(points_list, bit_num):
    phaseBit_PS_complex_list = []
    for points in points_list:
        phaseDeg_list = []
        for point in points:
            theta0 = point[0]
            phi0 = point[1]
            phaseRad = point_2_phase(theta0, phi0)
            phaseDeg = np.rad2deg(phaseRad)
            phaseDeg_list.append(phaseDeg)
        # _PS_complex
        phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_complex_beam_n(phaseDeg_list, bit_num)
        phaseBit_PS_complex_list.append(phaseBit_mix)
    return phaseBit_PS_complex_list


# ================================================== 流程相关 ==========================================================

# -------------------------------------------------- 统计不一样的位数 ------------------------------------------------
# 保存轨迹变化
def save_plots(list_data1, list_data2, list_data3, list_str, path_img_pre):
    # 检查列表长度是否一致
    if not (len(list_data1) == len(list_data2) == len(list_str)):
        raise ValueError("All input lists must have the same length.")
    # 逐图保存
    i = 1
    for i, (data1, data2, data3, str_label) in enumerate(zip(list_data1, list_data2, list_data3, list_str)):
        # 创建一个新的图像
        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(9, 3))
        # 在左上方绘制 list_data1[i]
        im1 = ax1.imshow(data1, aspect='auto')
        ax1.set_title(f"{str_label[0]}")
        ax1.axis('equal')
        ax1.axis('off')
        # plt.colorbar(im1, ax=ax1)     # 注释掉能关闭图例
        # 在中上方绘制 list_data2[i]
        im2 = ax2.imshow(data2, aspect='auto')
        ax2.set_title(f"{str_label[1]}")
        ax2.axis('equal')
        ax2.axis('off')
        # plt.colorbar(im2, ax=ax2)
        # 在右上方绘制 list_data3[i]
        im3 = ax3.imshow(data3, aspect='auto')
        ax3.set_title(f"{str_label[2]}")
        ax3.axis('equal')
        ax3.axis('off')
        # plt.colorbar(im3, ax=ax3)
        # 在下方放置文本
        fig.text(0.5, 0.01, str_label[3], ha='center', fontsize=12)
        # 调整子图间距
        plt.subplots_adjust(wspace=0.1, bottom=0.05)
        # 保存图像
        # plt.show()
        path_img = path_img_pre + str(i) + ".jpg"
        plt.savefig(path_img)
        # 关闭当前图像，释放内存
        plt.close(fig)
        i += 1

# 统计phases相邻phase[i]不相同的位数
def count_diff(phaseBits):
    counts = []
    for i in range(0, len(phaseBits)-1):
        phaseBit1 = phaseBits[i]
        phaseBit2 = phaseBits[i+1]
        # 比较 phase1 和 phase2 的每一位，值不相同则 count 加 1
        counts.append(np.sum(phaseBit1 != phaseBit2))
    return counts

# 每幅图保存: 相邻位置 - diff - count(diff)
def save_phases_diff(traces, phaseBits, counts, path_dir):
    list_data1 = []
    list_data2 = []
    list_data3 = []
    list_str = []
    for i in range(len(phaseBits) - 1):
        phase_this = phaseBits[i]
        phase_next = phaseBits[i + 1]
        diff_arr = compare_arrays(phase_this, phase_next)
        img_strs = [str(traces[i]), str(traces[i + 1]), "diff", "count diff:" + str(counts[i]) + ", total: 4096"]
        list_data1.append(phase_this)
        list_data2.append(phase_next)
        list_data3.append(diff_arr)
        list_str.append(img_strs)
    save_plots(list_data1, list_data2, list_data3, list_str, path_dir)

# 每幅图保存: phase - pattern - diff - count(diff)
def save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_dir):
    list_data1 = []
    list_data2 = []
    list_data3 = []
    list_str = []
    for i in range(len(phaseBits)):
        phaseBit = phaseBits[i]
        # fft计算方向图
        phaseDeg = phaseBit_2_deg(phaseBit, bit_num)
        phaseRad = np.deg2rad(phaseDeg)
        # pattern = phase_2_pattern(phaseBit_mix)                # 公式法直接计算, 准确但速度太慢
        pattern, x, y = phase_2_pattern_xyz_fft(phaseRad)  # FFT法计算, 快速
        # parse
        if i > 0:
            phaseBit_next = phaseBits[i-1]
            diff_arr = compare_arrays(phaseBit, phaseBit_next)
            img_strs = [str(traces[i]), "pattern", "diff", "count diff:" + str(counts[i - 1]) + ", total: 4096"]
        else:
            rows, cols = phaseBit.shape
            diff_arr = np.zeros((rows, cols))
            img_strs = [str(traces[i]), "pattern", "diff", ""]
        list_data1.append(phaseBit)
        list_data2.append(pattern)
        list_data3.append(diff_arr)
        list_str.append(img_strs)
    save_plots(list_data1, list_data2, list_data3, list_str, path_dir)


# =================================================== 主函数 ======================================================
# 沿(theta,phi)方向统计变化
def trace_counts_theta_phi_beam_2(theta1_start, theta1_step, theta1_end, phi1_start, phi1_step, phi1_end,
                                  theta2_start, theta2_step, theta2_end, phi2_start, phi2_step, phi2_end, path_pre):
    bit_num = 1
    traces = create_scan_traces_theta_phi_beam_2(theta1_start, theta1_step, theta1_end, phi1_start, phi1_step, phi1_end,
                                                 theta2_start, theta2_step, theta2_end, phi2_start, phi2_step, phi2_end)
    print("traces:")
    print(traces)
    phaseBits = trace_2_phaseBit_PS_complex(traces, bit_num)
    counts = count_diff(phaseBits)
    print("counts:", counts)
    # 每幅图保存: phase(t=i), phase(t=i+1), diff[phase(t=i+1) - phase(t=i)], count(diff)
    # save_phases_diff(traces, phaseBits, counts, path_pre + "compare_phases_diff/")
    # 每幅图保存: phase(t=i), pattern(t=i+1), diff[phase(t=i) - phase(t=i-1)], count(diff)
    save_phase_pattern_diff(bit_num, traces, phaseBits, counts, path_pre + "compare_phase_pattern_diff/")
    print("=======================================================================================================")


def main_beam_2():
    # 双波束: 一个波束不动, 一个沿theta方向变大
    # trace_counts_theta_phi_beam_2(45, 0, 45, 0, 0, 0, 0, 1, 90, 90, 0, 90,
    #                               "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(45,0,45)phi1(0,0,0)-theta2(0,1,90)phi2(90,0,90)/")
    trace_counts_theta_phi_beam_2(45, 0, 45, 0, 0, 0, 0, 1, 90, 180, 0, 180,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(45,0,45)phi1(0,0,0)-theta2(0,1,90)phi2(180,0,180)/")
    # 双波束: 都沿theta方向变大
    trace_counts_theta_phi_beam_2(0, 1, 90, 0, 0, 0, 0, 1, 90, 180, 0, 180,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(0,1,90)phi1(0,0,0)-theta2(0,1,90)phi2(180,0,180)/")
    trace_counts_theta_phi_beam_2(0, 1, 90, 0, 0, 0, 0, 1, 90, 90, 0, 90,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(0,1,90)phi1(0,0,0)-theta2(0,1,90)phi2(90,0,90)/")
    trace_counts_theta_phi_beam_2(0, 1, 90, 0, 0, 0, 0, 1, 90, 45, 0, 45,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(0,1,90)phi1(0,0,0)-theta2(0,1,90)phi2(45,0,45)/")
    #
    #
    # 双波束: 一个波束不动, 一个沿phi方向变大
    trace_counts_theta_phi_beam_2(45, 0, 45, 0, 0, 0, 30, 0, 30, 0, 1, 360,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(45,0,45)phi1(0,0,0)-theta2(30,0,30)phi2(0,1,360)/")
    trace_counts_theta_phi_beam_2(45, 0, 45, 0, 0, 0, 60, 0, 60, 0, 1, 360,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(45,0,45)phi1(0,0,0)-theta2(60,0,60)phi2(0,1,360)/")
    # 双波束: 一个沿phi方向顺时针, 一个沿phi逆时针
    trace_counts_theta_phi_beam_2(30, 0, 30, 0, 1, 360, 30, 0, 30, 360, -1, 0,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(30,0,30)phi1(0,1,360)-theta2(30,0,30)phi2(360,-1,0)/")
    # 双波束: 都沿phi逆时针
    trace_counts_theta_phi_beam_2(30, 0, 30, 0, 1, 180, 30, 0, 30, 180, 1, 360,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(30,0,30)phi1(0,1,180)-theta2(30,0,30)phi2(180,1,360)/")
    trace_counts_theta_phi_beam_2(30, 0, 30, 0, 1, 360, 60, 0, 60, 0, 1, 360,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(30,0,30)phi1(0,1,360)-theta2(60,0,60)phi2(0,1,360)/")
    #
    #
    # 双波束: 都沿(theta, phi)方向变大
    trace_counts_theta_phi_beam_2(10, 1, 70, 10, 1, 70, 10, 1, 70, 90, 1, 150,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(10,1,70)phi1(10,1,70)-theta2(10,1,70)phi2(90,1,150)/")
    trace_counts_theta_phi_beam_2(10, 1, 70, 10, 1, 70, 10, 1, 70, 180, 1, 240,
                                  "../files/multi-beam-trace/beam2/PS-complex/64x64/theta1(10,1,70)phi1(10,1,70)-theta2(10,1,70)phi2(180,1,240)/")




if __name__=="__main__":
    main_beam_2()